#ifndef __HTHREADPOOL_QUEUE_WORKSTEALINGQUEUE_H__
#define __HTHREADPOOL_QUEUE_WORKSTEALINGQUEUE_H__

#include <deque>
#include <mutex>
#include <queue>
#include <thread>
#include "CommonObject.h"
#include "QueueObject.h"
#include "Task.h"

HTHREADPOOL_NAMESPACE_BEGIN
    class WorkStealingQueue : public QueueObject
    {
    public:
        WorkStealingQueue() = default;
        THREADPOOL_NO_ALLOWED_COPY(WorkStealingQueue)
        HVoid push(Task&& task)
        {
            while(true) {
                if(m_mutex.try_lock()) {
                    m_deque.emplace_front(std::move(task));
                    m_mutex.unlock();
                    break;
                } else {
                    std::this_thread::yield();
                }
            }
        }

        HBool tryPop(Task& task)
        {
            // 这里不使用raii锁，主要是考虑到多线程的情况下，可能会重复进入
            bool result = false;
            if(m_mutex.try_lock()) {
                if(!m_deque.empty()) {
                    task = std::move(m_deque.front());
                    m_deque.pop_front();
                    result = true;
                }
                m_mutex.unlock();
            }

            return result;
        }

        HBool tryPop(TaskArrRef taskArr, int maxLocalBatchSize)
        {
            bool result = false;
            if(m_mutex.try_lock()) {
                while(!m_deque.empty() && maxLocalBatchSize--) {
                    taskArr.emplace_back(std::move(m_deque.front()));
                    m_deque.pop_front();
                    result = true;
                }
                m_mutex.unlock();
            }

            return result;
        }

        HBool trySteal(Task& task)
        {
            bool result = false;
            if(m_mutex.try_lock()) {
                if(!m_deque.empty()) {
                    task = std::move(m_deque.back());  // 从后方窃取
                    m_deque.pop_back();
                    result = true;
                }
                m_mutex.unlock();
            }

            return result;
        }

        HBool trySteal(TaskArrRef taskArr, int maxStealBatchSize)
        {
            bool result = false;
            if(m_mutex.try_lock()) {
                while(!m_deque.empty() && maxStealBatchSize--) {
                    taskArr.emplace_back(std::move(m_deque.back()));
                    m_deque.pop_back();
                    result = true;
                }
                m_mutex.unlock();
            }

            return result;  // 如果非空，表示盗取成功
        }

    private:
        std::deque<Task> m_deque;
    };
HTHREADPOOL_NAMESPACE_END

#endif
